package com.stress.testing.plugin.apache.rocketmq.interceptor;

import com.stress.testing.common.StressTesting;
import com.stress.testing.common.TraceContext;
import com.stress.testing.common.TraceContextHolder;
import com.stress.testing.common.interceptor.AbstractInterceptor;
import org.apache.rocketmq.common.MixAll;
import org.apache.rocketmq.common.message.MessageExt;

import java.lang.reflect.Method;
import java.util.Map;

import static com.stress.testing.common.conf.PluginConfigKey.DEFAULT_PLUGIN_APACHE_ROCKETMQ_TOPIC_PREFIX;
import static com.stress.testing.common.conf.PluginConfigKey.PLUGIN_APACHE_ROCKETMQ_TOPIC_PREFIX_KEY;

public class ProducerImplCheckTransactionStateInterceptor extends AbstractInterceptor {
    @Override
    public void before(Method method, Object[] args, Object zuper, Object thiz) {
        TraceContextHolder.removeTraceContext();

        if (!StressTesting.enable(getAgentConfig(), getPluginConfig())) {
            return;
        }

        if (args == null || args.length != 3) {
            return;
        }

        Object arg = args[1];
        if (arg == null || !(arg instanceof MessageExt)) {
            return;
        }

        MessageExt msgExt = (MessageExt) arg;
        String topic = msgExt.getTopic();
        String topicPrefix = getPluginConfig().get(PLUGIN_APACHE_ROCKETMQ_TOPIC_PREFIX_KEY,
                DEFAULT_PLUGIN_APACHE_ROCKETMQ_TOPIC_PREFIX);

        if (!topic.startsWith(topicPrefix) &&
                !topic.startsWith(MixAll.DLQ_GROUP_TOPIC_PREFIX + topicPrefix) &&
                !topic.startsWith(MixAll.RETRY_GROUP_TOPIC_PREFIX + topicPrefix)) {
            return;
        }

        TraceContext traceContext = new TraceContext(getAgentConfig().getTracePropertyPrefix());
        Map<String, String> properties = msgExt.getProperties();
        for (Map.Entry<String, String> entry : properties.entrySet()) {
            if (entry.getKey().startsWith(getAgentConfig().getTracePropertyPrefix())) {
                traceContext.setProperty(entry.getKey(), entry.getValue());
            }
        }

        // 可以简单的认为，如果msg的property中不包含trace_id,那么就不包含trace信息
        if (!traceContext.hasTraceID()) {
            return;
        }

        TraceContextHolder.setTraceContext(traceContext);
    }

    @Override
    public void afterReturning(Object result, Exception throwable) {
        TraceContextHolder.removeTraceContext();
    }
}
